import sqlite3
import os
import datetime
from typing import Dict, Any, List, Tuple

class Database:
    def __init__(self):
        """初始化数据库连接"""
        self.base_dir = os.path.dirname(os.path.abspath(__file__))
        self.db_path = os.path.join(self.base_dir, 'data', 'predictions.db')
        os.makedirs(os.path.dirname(self.db_path), exist_ok=True)
        self.connection = None
        self.cursor = None
        
    def connect(self) -> None:
        """连接到数据库"""
        self.connection = sqlite3.connect(self.db_path, check_same_thread=False)
        self.cursor = self.connection.cursor()
        self._create_tables()
        
    def close(self) -> None:
        """关闭数据库连接"""
        if self.connection:
            self.connection.close()
            self.connection = None
            self.cursor = None
            
    def _create_tables(self) -> None:
        """创建必要的数据库表"""
        # 创建预测结果表
        self.cursor.execute('''
        CREATE TABLE IF NOT EXISTS predictions (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            user_id TEXT,
            prediction INTEGER,
            probability REAL,
            timestamp TEXT,
            age INTEGER,
            sex INTEGER,
            cp INTEGER,
            trestbps INTEGER,
            chol INTEGER,
            fbs INTEGER,
            restecg INTEGER,
            thalach INTEGER,
            exang INTEGER,
            oldpeak REAL,
            slope INTEGER,
            ca INTEGER,
            thal INTEGER
        )
        ''')
        
        # 创建模型性能表
        self.cursor.execute('''
        CREATE TABLE IF NOT EXISTS model_performance (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            model_name TEXT,
            accuracy REAL,
            precision REAL,
            recall REAL,
            f1_score REAL,
            roc_auc REAL,
            avg_precision REAL,
            timestamp TEXT
        )
        ''')
        
        self.connection.commit()
        
    def _ensure_connection(self) -> None:
        """确保数据库连接存在"""
        if self.connection is None:
            self.connect()
        
    def save_prediction(self, prediction: int, probability: float, input_data: Dict[str, Any], user_id: str = None) -> int:
        """保存预测结果到数据库
        
        Args:
            prediction: 预测的类别（0或1）
            probability: 预测的概率值
            input_data: 输入的特征数据
            user_id: 用户提供的ID，默认为None
            
        Returns:
            新插入记录的ID
        """
        self._ensure_connection()
        timestamp = datetime.datetime.now().isoformat()
        
        query = '''
        INSERT INTO predictions 
        (user_id, prediction, probability, timestamp, age, sex, cp, trestbps, chol, 
         fbs, restecg, thalach, exang, oldpeak, slope, ca, thal)
        VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
        '''
        
        params = (
            user_id,
            prediction, 
            probability, 
            timestamp,
            input_data.get('age', 0),
            input_data.get('sex', 0),
            input_data.get('cp', 0),
            input_data.get('trestbps', 0),
            input_data.get('chol', 0),
            input_data.get('fbs', 0),
            input_data.get('restecg', 0),
            input_data.get('thalach', 0),
            input_data.get('exang', 0),
            input_data.get('oldpeak', 0.0),
            input_data.get('slope', 0),
            input_data.get('ca', 0),
            input_data.get('thal', 0)
        )
        
        self.cursor.execute(query, params)
        self.connection.commit()
        return self.cursor.lastrowid
        
    def get_recent_predictions(self, limit: int = 10) -> List[Tuple]:
        """获取最近的预测结果
        
        Args:
            limit: 返回的记录数量上限
            
        Returns:
            预测记录列表
        """
        self._ensure_connection()
        query = '''
        SELECT id, prediction, probability, timestamp 
        FROM predictions 
        ORDER BY id DESC 
        LIMIT ?
        '''
        
        self.cursor.execute(query, (limit,))
        return self.cursor.fetchall()
        
    def save_model_performance(self, model_name: str, metrics: Dict[str, float]) -> int:
        """保存模型性能指标
        
        Args:
            model_name: 模型名称
            metrics: 性能指标字典
            
        Returns:
            新插入记录的ID
        """
        self._ensure_connection()
        timestamp = datetime.datetime.now().isoformat()
        
        query = '''
        INSERT INTO model_performance 
        (model_name, accuracy, precision, recall, f1_score, roc_auc, avg_precision, timestamp)
        VALUES (?, ?, ?, ?, ?, ?, ?, ?)
        '''
        
        params = (
            model_name,
            metrics.get('accuracy', 0.0),
            metrics.get('precision', 0.0),
            metrics.get('recall', 0.0),
            metrics.get('f1', 0.0),
            metrics.get('roc_auc', 0.0),
            metrics.get('avg_precision', 0.0),
            timestamp
        )
        
        self.cursor.execute(query, params)
        self.connection.commit()
        return self.cursor.lastrowid
        
    def get_model_performance(self) -> List[Tuple]:
        """获取所有模型性能记录
        
        Returns:
            模型性能记录列表
        """
        self._ensure_connection()
        query = '''
        SELECT model_name, accuracy, precision, recall, f1_score, roc_auc, timestamp 
        FROM model_performance 
        ORDER BY timestamp DESC
        '''
        
        self.cursor.execute(query)
        return self.cursor.fetchall()
        
    def get_prediction_by_user_id(self, user_id: str) -> List[Tuple]:
        """根据用户ID获取预测记录
        
        Args:
            user_id: 用户ID
            
        Returns:
            匹配的预测记录列表
        """
        self._ensure_connection()
        
        query = '''
        SELECT id, user_id, prediction, probability, timestamp, 
               age, sex, cp, trestbps, chol, fbs, restecg, 
               thalach, exang, oldpeak, slope, ca, thal
        FROM predictions 
        WHERE user_id = ?
        ORDER BY timestamp DESC
        '''
        
        try:
            self.cursor.execute(query, (user_id,))
            results = self.cursor.fetchall()
            if results:
                print(f"获取了用户ID为 {user_id} 的 {len(results)} 条预测记录")
            else:
                print(f"未找到用户ID为 {user_id} 的预测记录")
            return results
        except Exception as e:
            print(f"获取用户预测记录错误: {str(e)}")
            raise

# 创建全局数据库实例
db = Database() 